import os
import logging
import time
from tqdm import tqdm
from magic_pdf.data.data_reader_writer import FileBasedDataWriter, FileBasedDataReader
from magic_pdf.data.dataset import PymuDocDataset
from magic_pdf.model.doc_analyze_by_custom_model import doc_analyze
from TableStructureRec.get_table_image import save_poly_image_from_pdf
from TableStructureRec.tableimg2html import batch_infer_table_structure
from utils.table_merge.table_md_merge import remove_body_html_tags
from utils.title_and_table_change import title_and_table_change
from log_config import setup_logging

# 日志初始化
setup_logging()


def pdf_to_markdown(
    pdf_path: str,
    output_dir: str,
    pre_md_save: bool = False,
    ocr: bool = False,
):
    """
    用TableStructureRec流水线，从PDF提取表格，识别结构，输出Markdown文件，整体流程自动化。
    Args:
        pdf_path (str): PDF文件路径
        output_dir (str): 输出文件夹路径
        pre_md_save (bool): 是否保存预处理的Markdown文件
        ocr (bool): 是否使用OCR进行文本识别
    Returns:
        final_md_path (str): 最终生成的Markdown文件路径
        md_content (str): 最终Markdown内容
    """
    start_time = time.time()
    # 路径与文件名准备
    name_without_extension = os.path.basename(pdf_path).split('.')[0]
    logging.info(f'运行{pdf_path:}')
    local_md_dir = output_dir + f"/{name_without_extension}"
    local_image_dir = output_dir + f"/{name_without_extension}/images"
    image_dir = str(os.path.basename(local_image_dir))
    os.makedirs(local_image_dir, exist_ok=True)

    image_writer = FileBasedDataWriter(local_image_dir)
    md_writer = FileBasedDataWriter(local_md_dir)

    # 读取PDF文件内容
    reader1 = FileBasedDataReader("")
    pdf_bytes = reader1.read(pdf_path)  # 读取PDF内容

    # 创建数据集实例并执行初步结构分析
    ds = PymuDocDataset(pdf_bytes)
    if ocr:
        # 如果需要OCR，执行OCR识别
        infer_result = ds.apply(doc_analyze, ocr=True)
    else:
        # 否则直接执行文档分析
        infer_result = ds.apply(doc_analyze, ocr=False)

    def table_predict_optimized(infer_result, pdf_path, local_image_dir):
        """
        优化表格识别流程：收集任务、批量截图、批量结构推理
        """
        # 1. 收集表格检测任务
        tasks = []
        for page in infer_result._infer_res:
            page_no = page['page_info']['page_no']
            page_size = (
                page['page_info']['width'],
                page['page_info']['height']
            )
            for layout_det in page['layout_dets']:
                if layout_det.get('category_id') == 5:  # 5表示表格
                    tasks.append({
                        'layout_det': layout_det,
                        'page_no': page_no,
                        'page_size': page_size,
                    })

        # 2. 批量保存表格图片
        for task in tqdm(tasks, desc="提取表格图片"):
            det = task['layout_det']
            img_path, _ = save_poly_image_from_pdf(
                pdf_path,
                task['page_no'],
                det['poly'],
                task['page_size'],
                local_image_dir
            )
            task['img_path'] = img_path

        # 3. 批量推理表格结构
        img_paths = [task['img_path'] for task in tasks]
        import time
        start_time = time.time()
        table_results = batch_infer_table_structure(
            img_paths=img_paths,
            model_type='auto',
            save=False
        )
        inference_time = time.time() - start_time
        for table_result, task in tqdm(zip(table_results, tasks), desc="表格识别", postfix={"总推理时间(s)": f"{inference_time:.2f}"}):
            det = task['layout_det']
            det['html'] = table_result.pred_html

        return infer_result

    # 表格识别主流程
    logging.info("开始进行自定义表格识别！！！")
    infer_result = table_predict_optimized(
        infer_result,
        pdf_path=pdf_path,
        local_image_dir=local_image_dir
    )

    # 生成Markdown内容，写入预处理md文件
    if ocr:
        # 如果使用OCR，使用OCR模式生成Markdown
        pipe_result = infer_result.pipe_ocr_mode(image_writer)
    else:
        # 否则使用文本模式生成Markdown
        pipe_result = infer_result.pipe_txt_mode(image_writer)
    # 保持处理前数据
    if pre_md_save:
        md_content = pipe_result.get_markdown(local_image_dir)
        md_content = remove_body_html_tags(md_content)
        pre_md_path = os.path.join(local_md_dir, f"{name_without_extension}_pre.md")
        md_writer.write_string(pre_md_path, md_content)

    # 尝试进行标题与表格样式调整
    try:
        md_content = title_and_table_change(pipe_result, local_image_dir, image_dir)
        logging.info("title_and_table_change 执行成功。")
    except Exception as e:
        logging.error(f"title_and_table_change 执行出错: {e}", exc_info=True)
        md_content = '程序错误，转换失败！！！'


    # 写入最终Markdown文件
    md_writer.write_string(f"{name_without_extension}.md", md_content)
    end_time = time.time()
    logging.info(f"Markdown文件保存成功\n{output_dir}/{name_without_extension}/{name_without_extension}.md")
    logging.info(f"all time:{end_time-start_time}")
    final_md_path = os.path.join(local_md_dir, f"{name_without_extension}.md")
    return final_md_path, md_content

if __name__ == "__main__":
    pdf_path = "..input_doc/sample/提取自德州德达城市建设投资运营有限公司2023年面向专业投资者公开发行公司债券（第一期）募集说明书(2).pdf"
    pdf_md_path, md_text = pdf_to_markdown(
        pdf_path=pdf_path,
        output_dir="../output"
    )
    print(f"Markdown文件路径: {pdf_md_path}")
